def maxProfit(k: int, prices) -> int:
    # dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k-1][1] + prices[i])
    # dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k][0] - prices[i])
    n = len(prices)
    if n <= 1:
        return 0
    #  k太大就失去约束作用：
    if k > n // 2:
        dp_0, dp_1 = 0, float('-inf')
        for i in range(n):
            tmp = dp_0
            dp_0 = max(dp_0, dp_1 + prices[i])
            dp_1 = max(dp_1, tmp - prices[i])
        return dp_0
    dp = [[[0] * 2 for _ in range(k + 1)] for _ in range(n + 1)]
    # 处理base条件：
    # 一开始不该有股票
    for i in range(k + 1):
        dp[0][i][1] = float('-inf')
    # 在不买入操作时不该有股票
    for i in range(n + 1):
        dp[i][0][1] = float('-inf')
    # print(dp)
    for i in range(1, n + 1):
        for j in range(k, 0, -1):
            dp[i][j][0] = max(dp[i - 1][j][1] + prices[i - 1], dp[i - 1][j][0])
            dp[i][j][1] = max(dp[i - 1][j][1], dp[i - 1][j - 1][0] - prices[i - 1])
    # print(dp)
    return dp[n][k][0]


def maxProfit_no_i(k: int, prices) -> int:
    # dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k-1][1] + prices[i])
    # dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k][0] - prices[i])
    n = len(prices)
    if n <= 1:
        return 0
    #  k太大就失去约束作用：
    if k > n // 2:
        dp_0, dp_1 = 0, float('-inf')
        for i in range(n):
            tmp = dp_0
            dp_0 = max(dp_0, dp_1 + prices[i])
            dp_1 = max(dp_1, tmp - prices[i])
        return dp_0
    dp_0 = [0] * (k + 1)
    dp_1 = [float('-inf')] * (k + 1)
    # 处理base条件：
    # print(dp)
    for i in range(n):
        tmp = list(dp_0)
        for j in range(k, 0, -1):

            dp_0[j] = max(dp_0[j], dp_1[j] + prices[i])
            dp_1[j] = max(dp_1[j], tmp[j - 1] - prices[i])
        # print(dp_0, dp_1)
    # print(dp)
    return dp_0[k]


if __name__ == '__main__':
    print(maxProfit(2, [3, 2, 6, 5, 0, 3]))
    print(maxProfit_no_i(2, [3, 2, 6, 5, 0, 3]))